package com.thrift.client.annotation;

import com.thrift.client.config.TTSocket;
import com.thrift.client.config.ThriftClientConnectPoolFactory;
import com.thrift.client.utils.SpringUtil;
import com.thrift.client.utils.ZkUtil;
import lombok.extern.slf4j.Slf4j;
import org.apache.commons.lang.StringUtils;
import org.apache.commons.pool.impl.GenericObjectPool;
import org.apache.curator.framework.CuratorFramework;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import org.springframework.beans.BeansException;
import org.springframework.beans.factory.config.BeanPostProcessor;
import org.springframework.context.ApplicationContext;
import org.springframework.stereotype.Component;

import javax.annotation.Resource;
import java.lang.reflect.Field;
import java.util.ArrayList;
import java.util.HashMap;
import java.util.List;
import java.util.Map;
import java.util.concurrent.ConcurrentHashMap;
import java.util.concurrent.atomic.AtomicInteger;

/**
 * @author 何锦涛--CodeSleep
 * @since 2022/9/4 20:27
 */
@Component
@Slf4j
public class ThriftClientBeanPostProcessor implements BeanPostProcessor {

    // RPC 服务连接池
    public static final Map<String, List<ThriftClientConnectPoolFactory>> connectPoolListMap = new ConcurrentHashMap<>();
    private static final Logger logger = LoggerFactory.getLogger(ThriftClientProxy.class);
    // RPC 所有相关节点公用的顶级父节点
    private static final String THRIFT_SERVER_PREFIX = "/thrift/";
    private static final AtomicInteger position = new AtomicInteger(0);

    @Resource
    private ThriftClientProxy thriftClientProxy;
    @Resource
    private final ApplicationContext applicationContext = SpringUtil.getApplicationContext();

    @Override
    public Object postProcessBeforeInitialization(Object bean, String beanName) throws BeansException {
        // 扫描当前类是否有客户端注解，如果有就创建 Thrift 客户端代理
        Field[] fields = bean.getClass().getDeclaredFields();
        for (Field field : fields) {
            if (field.getAnnotation(ThriftClient.class) != null) {
                try {
                    Object target = applicationContext.getBean(beanName);
                    Object clientProxy = createClient(bean, field.getType(), field.getAnnotation(ThriftClient.class));
                    field.set(target, clientProxy);
                } catch (Exception e) {
                    e.printStackTrace();
                    log.info("创建 Thrift 客户端代理错误：{}", field.getName());
                }
            }
        }
        return bean;
    }

    private Object createClient(Object bean, Class clazz, ThriftClient anno) throws Exception{
        // 创建 Curator 的 zk 连接
        CuratorFramework curatorFramework = ZkUtil.zookeeperClient();
        String host = anno.host();
        int port = anno.port();
        String serviceName = anno.serviceName();
        String key = "";
        // 如果没有 host 和 port，就通过 ServiceName 在 zk 中找，如果存在，获取里面的 host:port
        if (StringUtils.isBlank(host) || port == 0) {
            if (!StringUtils.isBlank(serviceName)) {
                if (ZkUtil.exist(THRIFT_SERVER_PREFIX.concat(serviceName))) {
                    List<String> children = ZkUtil.getChildren(THRIFT_SERVER_PREFIX.concat(serviceName));
                    if (children == null || children.size() == 0) {
                        logger.error("Thrift Server 未注册 zk 节点");
                    } else {
                        key = children.get(0);
                        host = StringUtils.substringBefore(key, ":");
                        port = Integer.parseInt(StringUtils.substringAfter(key, ":"));
                    }
                } else {
                    logger.error("Thrift Server 未注册 zk 节点，不存在 {} 节点", THRIFT_SERVER_PREFIX.concat(serviceName));
                    throw new RuntimeException("Thrift Server 未注册 zk 节点，不存在 "
                            + THRIFT_SERVER_PREFIX.concat(serviceName) + " 节点");
                }
            } else {
                logger.error("Thrift 未填写 port:host 以及 ServiceName");
                throw new RuntimeException("Thrift 未填写 port:host 以及 ServiceName");
            }
        } else {
            key = host + port;
        }
        List<ThriftClientConnectPoolFactory> poolList = connectPoolListMap.get(serviceName);
        ThriftClientConnectPoolFactory poolFactory;
        if (poolList == null || poolList.size() == 0) {
            poolList = new ArrayList<>();
            GenericObjectPool.Config config = new GenericObjectPool.Config();
            config.maxActive = 3;
            config.maxWait = 3000L;
            // 创建一个池工厂对象
            poolFactory = new ThriftClientConnectPoolFactory(config, host, port, serviceName);
            poolList.add(poolFactory);
            connectPoolListMap.put(serviceName, poolList);
            logger.info("在后置处理器中关于 {} 的连接池已创建", key + serviceName);
        } else {
            poolFactory = getLeastConnPollFactory(poolList);
            logger.info("在后置处理器有多个服务端，关于 {} 的连接池已创建", key + serviceName);
        }
        TTSocket connect = poolFactory.getConnect();
        return thriftClientProxy.bind(bean, clazz, connect, serviceName, poolFactory);
    }

    public static ThriftClientConnectPoolFactory getLeastConnPollFactory(List<ThriftClientConnectPoolFactory> pollList) {
        ArrayList<ThriftClientConnectPoolFactory> newConnectList = new ArrayList<>(pollList);
        ArrayList<ThriftClientConnectPoolFactory> serverList = new ArrayList<>(pollList);
        HashMap<Double, Integer> weightMap = new HashMap<>();

        ThriftClientConnectPoolFactory poolFactory = newConnectList.get(0);
        double weight = poolFactory.getWeight();
        for (int i = 0; i < weight; i++)
            serverList.add(poolFactory);

        for (ThriftClientConnectPoolFactory connect : newConnectList) {
            double weightCur = connect.getWeight();
            for (int i = 0; i < weightCur; i++)
                serverList.add(connect);
            log.info("比较的权重值为：{} [ weight {} ] <---> {} [ weightCur {} ]",
                    poolFactory.host + poolFactory.port, weight,
                    connect.host + connect.port, weightCur);
            if (weightCur < weight) {
                poolFactory = connect;
                weight = weightCur;
                weightMap.put(weight, weightMap.getOrDefault(weight, 0) + 1);
            }
        }

        if (weightMap.get(weight) > 1) {
            // 有多个后端的 conns/weight 值同为最小的，那么对它们采用加权轮询算法
            if (position.get() > serverList.size())
                position.set(0);
            poolFactory = serverList.get(position.get());
            position.incrementAndGet();
            log.info("目前的加权轮训下标为：{}", position.get());
            return poolFactory;
        } else {
            // 无冲突情况，直接返回权重最小的即可
            return poolFactory;
        }
    }
}
